#include <elf.h>

#include "trap.h"

#include "file.h"
#include "log.h"
#include "string.h"
#include "syscallno.h"
#include "console.h"
#include "vm.h"
#include "proc.h"
#include "memlayout.h"
#include "kalloc.h"
#include "mmu.h"

static inline int check_elf_header(Elf64_Ehdr *h) {
    unsigned char *e = h->e_ident;
    return
        e[EI_MAG0] == ELFMAG0 &&
        e[EI_MAG1] == ELFMAG1 &&
        e[EI_MAG2] == ELFMAG2 &&
        e[EI_MAG3] == ELFMAG3 &&
        e[EI_CLASS] == ELFCLASS64 &&
        e[EI_DATA] == ELFDATA2LSB &&
        e[EI_VERSION] == EV_CURRENT &&
        // e[EI_OSABI] == ELFOSABI_NONE &&
        e[EI_ABIVERSION] == 0 &&

        h->e_type == ET_EXEC &&
        h->e_machine == EM_AARCH64 &&
        h->e_version == EV_CURRENT &&
        h->e_ehsize == sizeof(Elf64_Ehdr) &&
        h->e_phoff >= h->e_ehsize &&
        h->e_phentsize == sizeof(Elf64_Phdr);
}

static inline int check_program_header(Elf64_Phdr *ph) {
    return
        ph->p_align == PGSIZE &&
        0 <= ph->p_offset &&
        (ph->p_offset - ph->p_vaddr) % PGSIZE == 0 &&
        (ph->p_flags & PF_R) &&
        ph->p_filesz <= ph->p_memsz &&
        ph->p_vaddr + ph->p_memsz >= ph->p_vaddr;
}

static int map_segment(uint64_t *pgdir, struct inode *ip, Elf64_Phdr *ph) {
    uint64_t i = ph->p_vaddr % PGSIZE;
    uint64_t vm_start = ph->p_vaddr - i;
    uint64_t vm_end = i + ph->p_memsz;
    uint64_t f_start = ph->p_offset - i;
    uint64_t f_end = i + ph->p_filesz;

    uint64_t off = 0;
    for ( ; i < vm_end; i += PGSIZE - off) {
        off = i % PGSIZE;

        uint64_t *ptr;
        char *page;
        if ((ptr = pgdir_access(pgdir, vm_start, &page)) == 0)
            return -1;

        uint64_t pte = *ptr;
        pte |= PTE_USER_FLAGS;
        if (ph->p_flags & PF_W)
            pte |= PTE_RW;
        else
            pte |= PTE_RO;
        *ptr = pte;

        vm_start += PGSIZE;

        if (i < f_end) {
            uint64_t read_size = MIN(PGSIZE - off, f_end - i);
            if (readi(ip, page + off, f_start + i, read_size) != read_size)
                return -1;
        }
    }

    return vm_start / PGSIZE;
}

int execve(const char *path, int argc, char *const argv[], int envc, char *const envp[]) {
    uint64_t *pgdir = 0;

    /**
     * load program into memory.
     */

    // open executable file
    struct inode *ip;
    begin_op();
    if ((ip = namei((char*) path)) == 0) {
        end_op();
        return -1;
    }
    ilock(ip);

    // check file type
    struct stat st;
    stati(ip, &st);
    if (st.st_mode == T_FILE)
        goto error;

    // read ELF header
    Elf64_Ehdr h;
    if (readi(ip, (char*) &h, 0, sizeof(h)) != sizeof(h))
        goto error;
    if (!check_elf_header(&h))
        goto error;
    // cprintf("execve: entry=0x%p\n", h.e_entry);

    // read program header and load segments
    if ((pgdir = pgdir_init()) == 0)
        return -1;

    int vm_size = 0;
    for (uint64_t i = 0; i < h.e_phnum; i++) {
        Elf64_Phdr ph;
        uint64_t poff = h.e_phoff + i * sizeof(ph);
        if (readi(ip, (char*) &ph, poff, sizeof(ph)) != sizeof(ph))
            goto error;
        if (ph.p_type != PT_LOAD)
            continue;

        /*
        cprintf("execve: program header #%lld:\n", i);
        cprintf("- ph.p_type=0x%x\n", ph.p_type);
        cprintf("- ph.p_flags=0x%x\n", ph.p_flags);
        cprintf("- ph.p_offset=0x%p\n", ph.p_offset);
        cprintf("- ph.p_vaddr=0x%p\n", ph.p_vaddr);
        cprintf("- ph.p_paddr=0x%p\n", ph.p_paddr);
        cprintf("- ph.p_filesz=%llu\n", ph.p_filesz);
        cprintf("- ph.p_memsz=%llu\n", ph.p_memsz);
        cprintf("- ph.p_align=%lld\n", ph.p_align);
        */

        if (!check_program_header(&ph))
            goto error;

        int vm_top;
        if ((vm_top = map_segment(pgdir, ip, &ph)) < 0)
            goto error;
        if (vm_top > vm_size)
            vm_size = vm_top;
    }

    iunlockput(ip);
    end_op();

    uint64_t ustack, sp, vsp;
    pvm_create_stack(pgdir, vm_size, &ustack, &vsp);
    sp = (uint64_t) (ustack + PGSIZE);

#define TO_VIRTUAL(x) (vsp + (x) - ustack)
#define SP_MASK 0xfffffffffffffff0

    /**
     * push argument strings.
     */

    // FIXME: implement envp.

    // 1 for argv's terminating null pointer, 1 for envp's, 4 for auxv
    uint64_t arr[MAXARG + 1 + 1 + 4];
    memset(arr, 0, sizeof(arr));

    uint64_t args_size = (argc + 1 + 4 + 1) * sizeof(uint64_t);
    arr[argc + 5] = 0;  // auxv
    arr[argc + 4] = AT_NULL;
    arr[argc + 3] = PGSIZE;
    arr[argc + 2] = AT_PAGESZ;
    arr[argc + 1] = 0;  // envp
    arr[argc + 0] = 0;  // argv

    for (int i = 0; i < argc; i++) {
        int len = strlen(argv[i]);
        sp = (sp - len - 1) & SP_MASK;
        if (sp < ustack)
            goto error;

        memcpy((char *) sp, argv[i], len + 1);  // +1 for NUL
        arr[i] = TO_VIRTUAL(sp);
    }

    sp = (sp - args_size - sizeof(uint64_t)) & SP_MASK;
    if (sp <= ustack)
        goto error;

    *((uint64_t *) sp) = argc;
    memcpy((char *) (sp + sizeof(uint64_t)), arr, args_size);

    /**
     * prepare trapframe
     */

    struct proc *p = thisproc();
    memset(p->tf, 0, sizeof(struct trapframe));
    p->tf->pc = h.e_entry;
    p->tf->stack = TO_VIRTUAL(sp);
    p->sz = vsp + PGSIZE;
    strncpy(p->name, path, sizeof(p->name));
    // p->chan = 0;

    uint64_t *old_pgdir = p->pgdir;
    assert(old_pgdir != 0);

    p->pgdir = pgdir;

    uvm_switch(p);
    vm_free(old_pgdir, 0);

    return 0;

#undef TO_VIRTUAL

    /*
     *
     * The initial stack is like
     *
     *   +-------------+
     *   | auxv[o] = 0 |
     *   +-------------+
     *   |    ....     |
     *   +-------------+
     *   |   auxv[0]   |
     *   +-------------+
     *   | envp[m] = 0 |
     *   +-------------+
     *   |    ....     |
     *   +-------------+
     *   |   envp[0]   |
     *   +-------------+
     *   | argv[n] = 0 |  n == argc
     *   +-------------+
     *   |    ....     |
     *   +-------------+
     *   |   argv[0]   |
     *   +-------------+
     *   |    argc     |
     *   +-------------+  <== sp
     *
     * where argv[i], envp[j] are 8-byte pointers and auxv[k] are
     * called auxiliary vectors, which are used to transfer certain
     * kernel level information to the user processes.
     *
     * ## Example
     *
     * ```
     * sp -= 8; *(size_t *)sp = AT_NULL;
     * sp -= 8; *(size_t *)sp = PGSIZE;
     * sp -= 8; *(size_t *)sp = AT_PAGESZ;
     *
     * sp -= 8; *(size_t *)sp = 0;
     *
     * // envp here. Ignore it if your don't want to implement envp.
     *
     * sp -= 8; *(size_t *)sp = 0;
     *
     * // argv here.
     *
     * sp -= 8; *(size_t *)sp = argc;
     *
     * // Stack pointer must be aligned to 16B!
     *
     * thisproc()->tf->sp = sp;
     * ```
     *
     */

error:
    if (pgdir)
        vm_free(pgdir, 0);
    if (ip) {
        iunlockput(ip);
        end_op();
    }
    return -1;
}
